#include <iostream>
#include <cmath>
#include <cuComplex.h>  // 使用 cuComplex 库

// FFT kernel
__global__ void fft_kernel(cuComplex* data, int N) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < N) {
        cuComplex temp = make_cuComplex(0.0f, 0.0f);  // 使用 make_cuComplex 创建复数

        // 执行 FFT 核心计算
        for (int k = 0; k < N; k++) {
            float realPart = cosf(-2.0f * M_PI * idx * k / N);
            float imagPart = sinf(-2.0f * M_PI * idx * k / N);
            cuComplex twiddle = make_cuComplex(realPart, imagPart);  // 使用 make_cuComplex 创建复数

            // 对数据进行加权求和
            temp = cuCaddf(temp, cuCmulf(data[k], twiddle));  // 使用 cuComplex 函数
        }

        data[idx] = temp;  // 保存结果
    }
}

int main() {
    const int N = 8;
    cuComplex* h_data_1 = new cuComplex[N];
    cuComplex* h_data_2 = new cuComplex[N];

    // 初始化数据
    for (int i = 0; i < N; i++) {
        h_data_1[i] = make_cuComplex(1.0f, 0.0f);  // 使用 make_cuComplex 初始化复数
    }

    for (int i = 0; i < N; i++) {
        h_data_2[i] = h_data_1[i];
    }

    cuComplex* d_data_1;
    cuComplex* d_data_2;
    cudaMalloc((void**)&d_data_1, N * sizeof(cuComplex));
    cudaMalloc((void**)&d_data_2, N * sizeof(cuComplex));
    cudaMemcpy(d_data_1, h_data_1, N * sizeof(cuComplex), cudaMemcpyHostToDevice);
    cudaMemcpy(d_data_2, h_data_2, N * sizeof(cuComplex), cudaMemcpyHostToDevice);

    // 启动 FFT kernel
    fft_kernel<<<1, N>>>(d_data_1, N);
    fft_kernel<<<1, N>>>(d_data_2, N);

    cudaMemcpy(h_data_1, d_data_1, N * sizeof(cuComplex), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_data_2, d_data_2, N * sizeof(cuComplex), cudaMemcpyDeviceToHost);

    // 输出结果
    int i;
    for (i = 0; i < N; i++) {
        if (cuCrealf(h_data_1[i]) != cuCrealf(h_data_2[i]) || 
            cuCimagf(h_data_1[i]) != cuCimagf(h_data_2[i])) {
            break;
        }
        std::cout << "Data[" << i << "] = " << cuCrealf(h_data_1[i]) << " + " 
                  << cuCimagf(h_data_1[i]) << "i" << std::endl;
    }
    if (i != N) {
        std::cout << "Error detected at index " << i << "\n";
    }

    // 清理
    delete[] h_data_1, h_data_2;
    cudaFree(d_data_1);
    cudaFree(d_data_2);

    return 0;
}
